import json
import re
import requests
import base64
import os
from urllib.parse import urlparse

class OCRProcessor:
    """OCR结果处理器，支持从文件或URL读取OCR结果并转换为Markdown格式"""
    
    def __init__(self, table_html_to_md=False):
        """
        初始化OCR处理器
        
        Args:
            table_html_to_md (bool): 是否将表格HTML转换为Markdown格式，默认False表示保持HTML格式
        """
        self.table_html_to_md = table_html_to_md
    
    def is_url(self, path_or_url):
        """
        判断输入是URL还是文件路径
        
        Args:
            path_or_url (str): 文件路径或URL
            
        Returns:
            bool: True表示是URL，False表示是文件路径
        """
        try:
            result = urlparse(path_or_url)
            return all([result.scheme, result.netloc])
        except:
            return False
    
    def load_ocr_data(self, path_or_url):
        """
        从文件路径或URL加载OCR数据
        
        Args:
            path_or_url (str): 文件路径或URL
            
        Returns:
            dict: OCR数据字典
            
        Raises:
            FileNotFoundError: 文件不存在
            requests.RequestException: URL请求失败
            json.JSONDecodeError: JSON解析失败
        """
        if self.is_url(path_or_url):
            # 从URL加载
            response = requests.get(path_or_url)
            response.raise_for_status()
            return response.json()
        else:
            # 从本地文件加载
            if not os.path.exists(path_or_url):
                raise FileNotFoundError(f"文件不存在: {path_or_url}")
            
            with open(path_or_url, 'r', encoding='utf-8') as f:
                return json.load(f)
    
    def image_to_base64(self, image_path):
        """
        将图片文件转换为base64编码
        
        Args:
            image_path (str): 图片文件路径
            
        Returns:
            str: base64编码的图片数据
            
        Raises:
            FileNotFoundError: 图片文件不存在
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图片文件不存在: {image_path}")
        
        with open(image_path, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
            # 根据文件扩展名确定MIME类型
            ext = os.path.splitext(image_path)[1].lower()
            mime_type = {
                '.jpg': 'image/jpeg',
                '.jpeg': 'image/jpeg',
                '.png': 'image/png',
                '.gif': 'image/gif',
                '.bmp': 'image/bmp',
                '.webp': 'image/webp'
            }.get(ext, 'image/jpeg')
            
            return f"data:{mime_type};base64,{encoded_string}"
    
    def send_ocr_request(self, image_path, ocr_url="http://127.0.0.1:5000/ocr", prompt_type="prompt_layout_all_en"):
        """
        发送OCR请求并获取结果
        
        Args:
            image_path (str): 图片文件路径
            ocr_url (str): OCR服务URL
            prompt_type (str): 提示类型
            
        Returns:
            dict: OCR结果
            
        Raises:
            Exception: 请求失败
        """
        # 设置会话
        session = requests.Session()
        session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36',
            'Accept': 'application/json, text/plain, */*',
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
            'Accept-Encoding': 'gzip, deflate, br',
            'Connection': 'keep-alive',
            'Referer': 'http://localhost:5173/',
            'Origin': 'http://localhost:5173',
            'Content-Type': 'application/json'
        })
        
        # 转换图片为base64
        base64_image = self.image_to_base64(image_path)
        
        # 准备OCR请求数据
        ocr_data = {
            "image": base64_image,
            "prompt_type": prompt_type,
            "stream": False
        }
        
        # 发送OCR请求
        response = session.post(ocr_url, json=ocr_data, timeout=60)
        
        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"OCR请求失败，状态码: {response.status_code}, 响应: {response.text}")
    
    def html_table_to_markdown(self, html_table):
        """
        将HTML表格转换为Markdown表格
        
        Args:
            html_table (str): HTML表格字符串
            
        Returns:
            str: Markdown表格
        """
        # 使用正则表达式提取表格内容
        table_pattern = r'<table[^>]*>(.*?)</table>'
        match = re.search(table_pattern, html_table, re.DOTALL | re.IGNORECASE)
        if not match:
            return html_table
        
        table_content = match.group(1)
        
        # 提取表头
        header_pattern = r'<thead[^>]*>(.*?)</thead>'
        header_match = re.search(header_pattern, table_content, re.DOTALL | re.IGNORECASE)
        headers = []
        
        if header_match:
            header_content = header_match.group(1)
            # 提取表头行
            th_pattern = r'<th[^>]*>(.*?)</th>'
            th_matches = re.findall(th_pattern, header_content, re.DOTALL | re.IGNORECASE)
            for th in th_matches:
                # 清理HTML标签和换行符
                clean_th = re.sub(r'<br\s*/?>', ' ', th, flags=re.IGNORECASE)  # 用空格替换<br>标签
                clean_th = re.sub(r'<[^>]+>', '', clean_th)  # 移除其他HTML标签
                clean_th = re.sub(r'\s+', ' ', clean_th)  # 合并多个空格
                clean_th = clean_th.strip()
                headers.append(clean_th)
        
        # 提取表格主体
        body_pattern = r'<tbody[^>]*>(.*?)</tbody>'
        body_match = re.search(body_pattern, table_content, re.DOTALL | re.IGNORECASE)
        rows = []
        
        if body_match:
            body_content = body_match.group(1)
            # 提取表格行
            tr_pattern = r'<tr[^>]*>(.*?)</tr>'
            tr_matches = re.findall(tr_pattern, body_content, re.DOTALL | re.IGNORECASE)
            
            for tr in tr_matches:
                # 提取单元格
                td_pattern = r'<td[^>]*>(.*?)</td>'
                td_matches = re.findall(td_pattern, tr, re.DOTALL | re.IGNORECASE)
                row = []
                for td in td_matches:
                    # 清理HTML标签，保留换行
                    clean_td = re.sub(r'<br\s*/?>', ' ', td, flags=re.IGNORECASE)  # 用空格替换<br>标签
                    clean_td = re.sub(r'<[^>]+>', '', clean_td)  # 移除其他HTML标签
                    clean_td = re.sub(r'\s+', ' ', clean_td)  # 合并多个空格
                    clean_td = clean_td.strip()
                    row.append(clean_td)
                if row:
                    rows.append(row)
        
        # 构建Markdown表格
        if not headers and not rows:
            return "无法解析表格"
        
        # 如果没有表头，使用第一行作为表头
        if not headers and rows:
            headers = rows.pop(0)
        
        # 创建表头
        md_table = "| " + " | ".join(headers) + " |\n"
        
        # 创建分隔线
        separator = "| " + " | ".join(["---"] * len(headers)) + " |\n"
        md_table += separator
        
        # 添加表格行
        for row in rows:
            # 确保每行的列数与表头一致
            while len(row) < len(headers):
                row.append("")
            while len(row) > len(headers):
                row.pop()
            
            md_table += "| " + " | ".join(row) + " |\n"
        
        return md_table
    
    def ocr_to_markdown(self, ocr_data):
        """
        将OCR识别结果转换为Markdown格式
        
        Args:
            ocr_data: OCR识别结果的JSON字符串或字典
            
        Returns:
            str: Markdown格式的文本
        """
        
        # 如果是字符串，解析为字典
        if isinstance(ocr_data, str):
            try:
                data = json.loads(ocr_data)
            except json.JSONDecodeError:
                return "无效的JSON数据"
        else:
            data = ocr_data
        
        # 提取响应数据
        response_data = data.get('response', [])
        if isinstance(response_data, str):
            try:
                response_data = json.loads(response_data)
            except json.JSONDecodeError:
                return "无效的响应数据"
        
        # 按类别分类内容
        categorized_content = {
            'Title': [],
            'Caption': [],
            'Text': [],
            'Table': [],
            'Page-footer': [],
            'Other': []
        }
        
        # 分类处理
        for item in response_data:
            category = item.get('category', 'Other')
            if category in categorized_content:
                categorized_content[category].append(item)
            else:
                categorized_content['Other'].append(item)
        
        # 构建Markdown
        markdown_parts = []
        
        # 标题
        if categorized_content['Title']:
            title_text = categorized_content['Title'][0].get('text', '')
            markdown_parts.append(f"# {title_text}\n")
        
        # 表标题
        if categorized_content['Caption']:
            caption_text = categorized_content['Caption'][0].get('text', '')
            markdown_parts.append(f"**{caption_text}**\n")
        
        # 文本说明
        if categorized_content['Text']:
            for text_item in categorized_content['Text']:
                text_content = text_item.get('text', '')
                markdown_parts.append(f"{text_content}\n")
        
        # 表格处理
        if categorized_content['Table']:
            for table_item in categorized_content['Table']:
                table_html = table_item.get('text', '')
                if self.table_html_to_md:
                    # 将HTML表格转换为Markdown表格
                    table_md = self.html_table_to_markdown(table_html)
                    markdown_parts.append(table_md)
                else:
                    # 保持原始HTML格式
                    markdown_parts.append(table_html)
        
        # 页脚
        if categorized_content['Page-footer']:
            for footer_item in categorized_content['Page-footer']:
                footer_text = footer_item.get('text', '')
                markdown_parts.append(f"\n{footer_text}")
        
        # 其他内容
        if categorized_content['Other']:
            markdown_parts.append("\n## 其他内容\n")
            for item in categorized_content['Other']:
                category = item.get('category', 'Unknown')
                text = item.get('text', '')
                markdown_parts.append(f"**{category}:** {text}\n")
        
        return "\n".join(markdown_parts)
    
    def process(self, input_path_or_url, output_md_path, image_path=None, ocr_url="http://127.0.0.1:5000/ocr"):
        """
        处理OCR数据并转换为Markdown文件
        
        Args:
            input_path_or_url (str): 输入的JSON文件路径或URL，如果提供image_path则此参数被忽略
            output_md_path (str): 输出的Markdown文件路径
            image_path (str, optional): 图片文件路径，如果提供则直接发送OCR请求
            ocr_url (str): OCR服务URL
            
        Returns:
            str: 生成的Markdown内容
        """
        try:
            if image_path:
                # 直接发送OCR请求
                print(f"正在处理图片: {image_path}")
                ocr_data = self.send_ocr_request(image_path, ocr_url)
            else:
                # 从文件或URL加载OCR数据
                print(f"正在加载OCR数据: {input_path_or_url}")
                ocr_data = self.load_ocr_data(input_path_or_url)
            
            # 转换为Markdown
            print("正在转换为Markdown格式...")
            markdown_result = self.ocr_to_markdown(ocr_data)
            
            # 保存到文件
            with open(output_md_path, "w", encoding="utf-8") as f:
                f.write(markdown_result)
            
            print(f"Markdown文件已保存为 {output_md_path}")
            return markdown_result
            
        except Exception as e:
            print(f"处理过程中出现错误: {e}")
            raise

# 使用示例
if __name__ == "__main__":
    # 默认保持HTML格式
    processor = OCRProcessor()
    
    # 或者转换为Markdown格式
    # processor = OCRProcessor(table_html_to_md=True)
    
    # 方式1: 从JSON文件转换为Markdown
    # processor.process("response-nostream.json", "output_from_file.md")
    
    # 方式2: 从URL转换为Markdown
    # processor.process("http://example.com/ocr_result.json", "output_from_url.md")
    
    # 方式3: 直接处理图片并转换为Markdown
    # processor.process(None, "output_from_image.md", image_path="./demo/demo_image1.jpg")
    
    print("OCR处理器已准备就绪，可以通过以下方式使用：")
    print("1. processor.process('json_file_path', 'output.md')  # 从JSON文件处理")
    print("2. processor.process('http://example.com/result.json', 'output.md')  # 从URL处理")
    print("3. processor.process(None, 'output.md', image_path='image.jpg')  # 直接处理图片")
    print("\n注意: 初始化时设置table_html_to_md=True可将HTML表格转换为Markdown格式")